import cv2
import mediapipe as mp
import time
import pygame
import keyboard
import os

# --- 配置参数 ---
SITTING_TIME_LIMIT_MINUTES = 45  # 坐姿时间限制（分钟）
SITTING_TIME_LIMIT_SECONDS = SITTING_TIME_LIMIT_MINUTES * 60
# 用户需要将录制好的MP3文件命名为此，并放在脚本同目录下
AUDIO_FILENAME = "reminder_audio.mp3"
WINDOW_NAME = "Eye Keeper - Press ESC to quit"

# --- 主程序 ---
def main():
    audio_available = False
    # 1. 初始化 Pygame Mixer (用于播放音频)
    if os.path.exists(AUDIO_FILENAME):
        try:
            pygame.mixer.init()
            pygame.mixer.music.load(AUDIO_FILENAME)
            audio_available = True
            print(f"提示：已成功加载音频文件 '{AUDIO_FILENAME}'。")
        except pygame.error as e:
            print(f"错误：初始化Pygame Mixer或加载音频文件 '{AUDIO_FILENAME}' 失败 - {e}")
            print(f"请确保 '{AUDIO_FILENAME}' 是一个有效的MP3文件，并且与脚本在同一目录下。")
            audio_available = False
    else:
        print(f"警告：未找到音频文件 '{AUDIO_FILENAME}'。")
        print(f"请将您录制好的提醒MP3文件命名为 '{AUDIO_FILENAME}' 并放置在脚本同目录下。")
        print("提醒功能将无法播放声音。")

    # 2. 初始化 MediaPipe Pose
    mp_pose = mp.solutions.pose
    pose = mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5)
    mp_drawing = mp.solutions.drawing_utils

    # 3. 初始化摄像头
    cap = cv2.VideoCapture(0) # 0 代表默认摄像头
    if not cap.isOpened():
        print("错误：无法打开摄像头。请检查摄像头是否连接并可用。")
        if audio_available and pygame.mixer.get_init():
            pygame.mixer.quit()
        return

    # 4. 初始化状态变量
    person_present_timer_start = None  # 记录人开始出现的时间戳
    reminder_triggered_for_current_sitting = False # 标记当前久坐是否已触发提醒
    last_space_press_time = 0 # 用于空格键去抖

    print("\n程序启动成功。将摄像头对准您的学习区域。")
    print(f"当检测到您连续坐满 {SITTING_TIME_LIMIT_MINUTES} 分钟后，将播放语音提醒。")
    print("按【空格键】可以停止当前提醒并重置计时。")
    print("按【ESC键】退出程序。")
    if not audio_available:
        print("注意：由于音频文件问题，语音提醒功能当前不可用。")


    try:
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                print("错误：无法从摄像头读取帧。")
                break

            # 镜像翻转画面，使其更像镜子
            frame = cv2.flip(frame, 1)

            # MediaPipe 处理
            image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = pose.process(image_rgb)
            image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR) # 转回 BGR 用于 OpenCV 显示

            person_detected_this_frame = results.pose_landmarks is not None

            # 绘制姿态关键点 (可选)
            if person_detected_this_frame:
                mp_drawing.draw_landmarks(
                    image_bgr, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)

            # 状态更新与计时逻辑
            current_time = time.time()

            if person_detected_this_frame:
                if person_present_timer_start is None: # 人刚出现或中断后重新出现
                    person_present_timer_start = current_time
                    reminder_triggered_for_current_sitting = False # 为新的久坐周期重置提醒标记
                    print("检测到人体，开始计时...")
                
                elapsed_sitting_time = current_time - person_present_timer_start
                
                # 在画面上显示计时
                timer_text = f"You have been sitted for: {int(elapsed_sitting_time // 60)}m {int(elapsed_sitting_time % 60)}s"
                cv2.putText(image_bgr, timer_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

                # 检查是否达到提醒时间
                if audio_available and \
                   elapsed_sitting_time >= SITTING_TIME_LIMIT_SECONDS and \
                   not reminder_triggered_for_current_sitting and \
                   not pygame.mixer.music.get_busy():
                    print(f"已连续坐满 {SITTING_TIME_LIMIT_MINUTES} 分钟，播放提醒...")
                    pygame.mixer.music.play()
                    reminder_triggered_for_current_sitting = True
            
            else: # 当前帧未检测到人
                if person_present_timer_start is not None:
                    print("未检测到人体，计时暂停。")
                person_present_timer_start = None # 重置计时器起点
                # 注意：如果人离开时提醒正在播放，它会继续播放完或直到按空格

            # 处理空格键取消提醒 (带去抖)
            if keyboard.is_pressed('space'):
                if current_time - last_space_press_time > 0.5: # 0.5秒去抖
                    last_space_press_time = current_time
                    if audio_available and pygame.mixer.music.get_busy():
                        pygame.mixer.music.stop()
                        print("提醒已通过空格键停止。")
                    
                    print("计时器已重置。")
                    person_present_timer_start = current_time if person_detected_this_frame else None # 如果人还在，则从现在开始重新计时
                    reminder_triggered_for_current_sitting = False


            # 显示画面
            cv2.imshow(WINDOW_NAME, image_bgr)

            # 按 ESC 键退出
            if cv2.waitKey(5) & 0xFF == 27:
                print("ESC键按下，正在退出程序...")
                break
    finally:
        # 清理资源
        print("正在清理资源...")
        cap.release()
        if 'pose' in locals() and hasattr(pose, 'close'): # 确保 pose 已初始化
             pose.close()
        cv2.destroyAllWindows()
        if audio_available and pygame.mixer.get_init(): # 检查 pygame.mixer 是否已初始化
            pygame.mixer.music.stop() # 确保音乐停止
            pygame.mixer.quit()
        print("程序已退出。")

if __name__ == '__main__':
    main()